import json
import ollama
from PluginManager import PluginManager


class ChatBot:
    def __init__(self, modelname="qwen2.5:3b"):
        self.model = modelname
        self.messages = []
        self.messages.append({"role": "system", "content": """你是我创造的个人助理，你可以调用我的工具，来帮助我完成任务。
        所有问题请用中文回复。"""})
        self.tool_descriptions =[]
        self.pluginManager = PluginManager()
        self.pluginManager.load_plugins()
        # 生成Ollama需要的tools格式 
        for fn_Name, fn_Meta in self.pluginManager.tools.items(): 
            self.tool_descriptions.append(
                {
                    "type": "function",
                    "function": {
                        "name": fn_Name,
                        "parameters": {
                            "type": "object",
                            "required": None if len(fn_Meta["required"]) == 0 else fn_Meta["required"],
                            "properties": None if len(fn_Meta["parameters"]) == 0 else fn_Meta["parameters"]
                        }
                    }
                }
            )
             
    def chatBot(self, user_input: str):
        self.messages.append({"role": "user", "content": user_input})
        print(self.model + "---->" + user_input)
        # 在ollama.chat调用设置model, messages, tools参数
        response = ollama.chat(
            model=self.model, 
            messages=self.messages,
            tools=self.tool_descriptions,
            stream=False
        )   
        msg = response["message"]
        # print(msg)
        self.messages.append(msg)  # 保存原始响应
        # 处理工具调用
        if msg.get("tool_calls"):
            tool_results = self._ToolsCallFunction(msg["tool_calls"])
            # 添加工具执行结果
            self.messages.extend(tool_results)
            if len(tool_results) == 1:
                strcon = tool_results[0]["content"]
                return strcon
            else:
                return "Error: 未找到工具调用结果"
        else:
            # 直接返回文本响应
            if msg["content"].startswith("<tool_call"):
                return "有些卡壳啊，再说下。。。"  
            else:
                return msg["content"]  


    def _ToolsCallFunction(self, tool_calls):
        """处理多个工具调用
        """
        results = []
        for call in tool_calls:
            func_name =call.function.name; #call["function"]["name"]
            if func_name not in self.pluginManager.functions: 
                results.append(
                    {"role": "tool", "name": func_name, "content": f"Error: 未注册函数{func_name}"}
                )
                continue

            try:
                # 解析参数（需处理格式问题）
                objDict = call["function"]["arguments"]
                strs = json.dumps(objDict)
                args = json.loads(strs)
                func = self.pluginManager.functions[func_name]
                
                if len(func.__tool_meta__["parameters"]) != len(objDict):
                    results.append(
                        {"role": "tool", "name": func_name, "content": f"Error:{func_name}->参数无法正确匹配。"}
                    )
                    continue  
                result = func(**args)
                results.append(
                    {"role": "tool", "name": func_name, "content": str(result)}
                )
            except Exception as e: 
                results.append(
                    {"role": "tool", "name": func_name, "content": str(e)}
                )

        return results


    def chatBotStream(self, user_input: str):
        self.messages.append({"role": "user", "content": user_input})
        print(self.model + "---->" + user_input)
        # 在ollama.chat调用设置model, messages, tools参数
        response = ollama.chat(
            model=self.model, 
            messages=self.messages,
            tools=self.tool_descriptions,
            stream=True
        ) 
        streaming = True
                     
        if streaming:
            for part in response:
                if "message" in part:
                    msg = part["message"]
                    self.messages.append(msg)  # 保存原始响应
                    if msg.get("tool_calls"):
                        tool_results = self._ToolsCallFunction(msg["tool_calls"])
                        # 添加工具执行结果
                        self.messages.extend(tool_results)
                        if len(tool_results) == 1:
                            strcon = tool_results[0]["content"] 
                            print(strcon, end='', flush=True )
                        else: 
                            print("Error: 未找到工具调用结果", end='', flush=True )
                    else:
                        print(msg.content, end='', flush=True )
        else:  
            print( response['message']['content'], end='', flush=True )
